package statistics;

import java.util.Random;

/**
 * This class represents a Mix of Negative Binomial Distributions.
 */

public class MixNegBinomialDistribution extends DiscreteDistribution {

  private double p;
  private int  k;
  private double q;

  /**
   * Constructs a Mix of Negative Binomial Distributions with parameter k and p. You must also specify the
   * random seed that will be used to realize random variables with this distribution.
   * @param expect the expectation of the Mix of Negative Binomial Distributions
   * @param a
   * @param random the random seed.
   */

  public MixNegBinomialDistribution(double expect, double a, Random random) {
    this.random = random;
    k = ((int)Math.floor(1/a));
    q = ( a*(1 + k) - Math.sqrt((1-a*k)*(1 + k)))/(1 + a);
    p = expect/(k + 1 - q + expect);
  }

  /**
   * Returns the expectation of a Mix of negative binomial distributions with parameter k and p.
   */

  public double expectation() {
    return q*k/p + (1-q)*(k+1)/p;
  }

  /**
   * Returns the variance of a Mix of negative binomial distributions with parameter k and p.
   */

  public double variance() {
    return q*k*(1-p)/(p*p) + (1-q)*(k+1)*(1-p)/(p*p);
  }


  /**
   * Returns a realization of a random variable with a Mix of negative binomial distributions.
   */

  public double nextRandom() {
     double u1 = random.nextDouble();
     GeometricDistribution geo = new GeometricDistribution(p, random);
     double sum = 0;
     if(u1 < q) {
        for(int i = 0; i < k;i++) {
           double u = geo.nextRandom();
           sum = sum + u;
        }
     }
     else {
        for(int i = 0; i < k+1;i++) {
           double u = geo.nextRandom();
           sum = sum + u;
        }
     }
     return sum;
  }

  private double negBinom(int kk, double pp, int x) {
    double denom = 1.0;
    double nom = 1.0;
    // k * ... * (x+1)
    for (int i = x+kk-1; i > kk-1; i--) {
      denom = denom * i;
    }
    // (k-x)!
    for (int i = x; i > 1; i--) {
      nom = nom * i;
    }
    return denom/nom * Math.pow(1-pp, kk) * Math.pow(pp, x);
  }

  public double probabilityFunction(int x) {
    return q*negBinom(k, p, x)+(1.0-q)*negBinom(k+1, p, x);
  }
}
